#!/usr/bin/env python3
"""
本地LM模型测试脚本
"""

import requests
import json
import sys

def test_local_model_connection():
    """测试本地模型连接"""
    print("[SEARCH] 测试本地LM模型连接...")
    
    try:
        # 测试模型列表接口
        response = requests.get("http://localhost:1234/v1/models", timeout=10)
        if response.status_code == 200:
            models = response.json()
            print(f"[PASS] 本地LM服务连接正常")
            print(f"   可用模型: {[model['id'] for model in models.get('data', [])]}")
            return True
        else:
            print(f"[FAIL] 本地LM服务响应异常: {response.status_code}")
            return False
    except requests.RequestException as e:
        print(f"[FAIL] 本地LM服务连接失败: {e}")
        return False

def test_local_model_chat():
    """测试本地模型聊天功能"""
    print("[SEARCH] 测试本地LM模型聊天功能...")
    
    try:
        data = {
            "model": "google/gemma-3-4b",
            "messages": [
                {
                    "role": "system",
                    "content": "你是一个专业的日志分析专家，请用中文回答问题。"
                },
                {
                    "role": "user",
                    "content": "请简单介绍一下日志分析的重要性。"
                }
            ],
            "temperature": 0.7,
            "max_tokens": -1,
            "stream": False
        }
        
        response = requests.post(
            "http://localhost:1234/v1/chat/completions",
            headers={"Content-Type": "application/json"},
            json=data,
            timeout=30
        )
        
        if response.status_code == 200:
            result = response.json()
            answer = result['choices'][0]['message']['content']
            print("[PASS] 本地LM模型聊天功能正常")
            print(f"   测试回答: {answer[:100]}...")
            return True
        else:
            print(f"[FAIL] 本地LM模型聊天请求失败: {response.status_code}")
            print(f"   错误信息: {response.text}")
            return False
    except requests.RequestException as e:
        print(f"[FAIL] 本地LM模型聊天测试失败: {e}")
        return False

def test_log_analysis_prompt():
    """测试日志分析提示词"""
    print("[SEARCH] 测试日志分析提示词...")
    
    try:
        # 模拟日志分析数据
        log_data = {
            "total_logs": 1000,
            "error_count": 50,
            "warning_count": 100,
            "components": ["bmc_core", "bmc_network", "bmc_web"],
            "error_types": ["memory_error", "network_error", "permission_error"]
        }
        
        data = {
            "model": "google/gemma-3-4b",
            "messages": [
                {
                    "role": "system",
                    "content": "你是一个专业的日志分析专家，专门帮助用户分析系统日志中的问题。请用中文回答用户的问题。"
                },
                {
                    "role": "user",
                    "content": f"请分析以下日志数据并识别主要问题：\n{json.dumps(log_data, ensure_ascii=False, indent=2)}\n\n请提供：\n1. 主要问题概述\n2. 问题优先级排序\n3. 建议的解决方案"
                }
            ],
            "temperature": 0.7,
            "max_tokens": -1,
            "stream": False
        }
        
        response = requests.post(
            "http://localhost:1234/v1/chat/completions",
            headers={"Content-Type": "application/json"},
            json=data,
            timeout=60
        )
        
        if response.status_code == 200:
            result = response.json()
            answer = result['choices'][0]['message']['content']
            print("[PASS] 日志分析提示词测试正常")
            print(f"   分析结果: {answer[:200]}...")
            return True
        else:
            print(f"[FAIL] 日志分析提示词测试失败: {response.status_code}")
            return False
    except requests.RequestException as e:
        print(f"[FAIL] 日志分析提示词测试失败: {e}")
        return False

def main():
    """主测试函数"""
    print("[TEST] 本地LM模型测试")
    print("=" * 50)
    
    # 测试连接
    if not test_local_model_connection():
        print("\n[FAIL] 本地LM服务未启动或无法连接")
        print("   请确保LM Studio或其他本地模型服务正在运行")
        print("   默认地址: http://localhost:1234")
        sys.exit(1)
    
    # 测试聊天功能
    if not test_local_model_chat():
        print("\n[FAIL] 本地LM模型聊天功能异常")
        sys.exit(1)
    
    # 测试日志分析
    if not test_log_analysis_prompt():
        print("\n[FAIL] 日志分析提示词测试失败")
        sys.exit(1)
    
    print("\n" + "=" * 50)
    print("[SUCCESS] 本地LM模型测试完成!")
    print("\n[PASS] 所有测试通过")
    print("   本地LM模型已准备就绪")
    print("   可以在AI日志分析系统中使用")

if __name__ == "__main__":
    main() 